# -*- coding: utf-8 -*-
"""
@Time: 2025/6/19 23:45
@Author: luyinglin
@Email: 2902801287@qq.com
@File: travel-mysql.py
"""

import pandas as pd
from sqlalchemy import create_engine, Column, Integer, String, Float, Text, Boolean
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy.exc import IntegrityError, DataError
import traceback


# 1. 读取原始Excel数据
def load_and_preprocess_data(file_path):
    """加载Excel数据并进行预处理"""
    try:
        # 读取Excel文件
        df = pd.read_excel(file_path)
        print(f"原始数据量: {len(df)}条")

        # 数据清洗
        # 移除销量为0的记录
        df = df[df['销量'] != 0]
        # 填充空值
        df.fillna('未知', inplace=True)

        print(f"清洗后数据量: {len(df)}条")
        return df
    except Exception as e:
        print(f"数据加载失败: {str(e)}")
        return None


# 2. 省份映射函数
def map_province(df):
    """将省/市/区映射为完整省份名称"""
    # 中国省份和自治区的名称映射字典
    province_mapping = {
        '北京': '北京市',
        '天津': '天津市',
        '河北': '河北省',
        '山西': '山西省',
        '内蒙古': '内蒙古自治区',
        '辽宁': '辽宁省',
        '吉林': '吉林省',
        '黑龙江': '黑龙江省',
        '上海': '上海市',
        '江苏': '江苏省',
        '浙江': '浙江省',
        '安徽': '安徽省',
        '福建': '福建省',
        '江西': '江西省',
        '山东': '山东省',
        '河南': '河南省',
        '湖北': '湖北省',
        '湖南': '湖南省',
        '广东': '广东省',
        '广西': '广西壮族自治区',
        '海南': '海南省',
        '重庆': '重庆市',
        '四川': '四川省',
        '贵州': '贵州省',
        '云南': '云南省',
        '西藏': '西藏自治区',
        '陕西': '陕西省',
        '甘肃': '甘肃省',
        '青海': '青海省',
        '宁夏': '宁夏回族自治区',
        '新疆': '新疆维吾尔自治区',
        '台湾': '台湾省',
        '香港': '香港特别行政区',
        '澳门': '澳门特别行政区'
    }
    # 添加省份列
    df['省份'] = df['城市'].map(province_mapping).fillna('未知省份')
    return df


# 3. 数据库模型定义
Base = declarative_base()


class TravelAttraction(Base):
    """旅游景点数据库模型"""
    __tablename__ = 'attractions'

    id = Column(Integer, primary_key=True, autoincrement=True)
    city = Column(String(50), nullable=False, comment="城市")
    name = Column(String(100), nullable=False, index=True, comment="景点名称")
    star_level = Column(String(20), nullable=False, comment="星级")
    rating = Column(Float, nullable=True, comment="评分")
    price = Column(Float, nullable=True, comment="价格")
    sales = Column(Integer, nullable=True, comment="销量")
    region = Column(String(50), nullable=False, comment="省/市/区")
    coordinate = Column(String(50), nullable=True, comment="坐标")
    introduction = Column(Text, nullable=True, comment="简介")
    is_free = Column(Boolean, nullable=True, comment="是否免费")
    location = Column(String(200), nullable=True, comment="具体位置")
    province = Column(String(50), nullable=False, index=True, comment="所属省份")


# 4. 数据库操作类
class DatabaseManager:
    """数据库管理类"""

    def __init__(self, db_uri):
        self.engine = create_engine(
            db_uri,
            pool_size=5,
            pool_recycle=3600,
            connect_args={'connect_timeout': 10}
        )
        self.Session = sessionmaker(bind=self.engine)

    def create_tables(self):
        """创建数据库表"""
        try:
            Base.metadata.create_all(self.engine)
            print("数据库表创建成功")
            return True
        except Exception as e:
            print(f"数据库表创建失败: {str(e)}")
            return False

    def insert_data(self, df):
        """将DataFrame数据插入数据库"""
        session = self.Session()
        success_count = 0
        error_count = 0
        batch_size = 100

        print("开始插入数据...")

        for index in range(0, len(df), batch_size):
            batch = df.iloc[index:index + batch_size]
            try:
                for _, row in batch.iterrows():
                    try:
                        # 处理是否免费字段
                        is_free = None
                        if row['是否免费'] == '免费':
                            is_free = True
                        elif row['是否免费'] == '收费':
                            is_free = False

                        attraction = TravelAttraction(
                            city=str(row['城市'])[:50],
                            name=str(row['名称'])[:100],
                            star_level=str(row['星级'])[:20],
                            rating=float(row['评分']) if pd.notna(row['评分']) else None,
                            price=float(row['价格']) if pd.notna(row['价格']) else None,
                            sales=int(row['销量']) if pd.notna(row['销量']) else None,
                            region=str(row['省/市/区'])[:50],
                            coordinate=str(row['坐标'])[:50] if '坐标' in row and pd.notna(row['坐标']) else None,
                            introduction=str(row['简介'])[:5000] if '简介' in row and pd.notna(row['简介']) else None,
                            is_free=is_free,
                            location=str(row['具体位置'])[:200] if '具体位置' in row and pd.notna(
                                row['具体位置']) else None,
                            province=str(row['省份'])[:50]
                        )
                        session.add(attraction)
                        success_count += 1
                    except (ValueError, TypeError) as e:
                        error_count += 1
                        print(f"行{index}数据转换错误: {e}")
                        continue

                session.commit()
                print(f"已提交批次: {index // batch_size + 1}, 成功: {success_count}, 失败: {error_count}")

            except (IntegrityError, DataError) as e:
                session.rollback()
                print(f"批处理错误（行{index}-{index + batch_size}）:\n{traceback.format_exc()}")
                error_count += len(batch)

        session.close()
        print(f"数据插入完成! 总计: 成功 {success_count} 条, 失败 {error_count} 条")
        return success_count, error_count


# 5. 主函数
def main():
    # 文件路径
    input_file = r'城市景点信息表.xlsx'

    # 数据库连接信息
    DB_URI = 'mysql+pymysql://root:root@localhost/stat'

    # 加载数据
    df = load_and_preprocess_data(input_file)
    if df is None:
        print("数据处理失败，程序退出")
        return

    # 映射省份
    df = map_province(df)


    # 数据库操作
    db_manager = DatabaseManager(DB_URI)

    # 创建表
    if not db_manager.create_tables():
        print("数据库表创建失败，程序退出")
        return

    # 插入数据
    success, errors = db_manager.insert_data(df)



if __name__ == "__main__":
    main()